import json
import time
from dotenv import load_dotenv
load_dotenv()
from zhipuai import ZhipuAI
from api import get_characterglm_response


def characterglm_assistant(messages,bot_name,bot_info, user_name, user_info):

    character_meta = {
        "user_info": user_info,
        "bot_info": bot_info,
        "user_name": user_name,
        "bot_name": bot_name
    }
    message = messages[-2:]
    # print(message)
    chunks = ""
    for chunk in get_characterglm_response(message, meta=character_meta):
        # print(chunk)
        # time.sleep(0.5)
        chunks = chunks + chunk
    print(chunks)
    messages.append({"role": "assistant", "content": chunks})
    return messages

def characterglm_user(messages,bot_name,bot_info, user_name, user_info):

    character_meta = {
        "user_info": user_info,
        "bot_info": bot_info,
        "user_name": user_name,
        "bot_name": bot_name
        # "bot_info": user_info,
        # "user_info": bot_info,
        # "bot_name": user_name,
        # "user_name": bot_name
    }
    message = messages[-2:]
    # print(message)
    chunks = ""
    for chunk in get_characterglm_response(messages, meta=character_meta):
        # print(chunk)
        # time.sleep(0.5)
        chunks = chunks + chunk
    print(chunks)
    messages.append({"role": "user", "content": chunks})
    return messages

def characterglm_example(bot_name,bot_info, user_name, user_info):

    character_meta = {
        "user_info": user_info,
        "bot_info": bot_info,
        "user_name": user_name,
        "bot_name": bot_name
    }
    messages = [
        {"role": "assistant", "content": "少安，我会死吗？"},
        {"role": "user", "content": "怎么会呢？再乱想我锤死你"},
        {"role": "assistant", "content": "那我就好好想想"}
    ]
    for chunk in get_characterglm_response(messages, meta=character_meta):
        print(chunk)
        time.sleep(0.5)



# 根据输入文本，调用glm-4，生成角色人设，并提取
def generate_role_play(text):
    client = ZhipuAI()
    content = f"请根据以下文本，生成两个角色的姓名和人设，返回值为json标准格式，关键字为姓名和人设，内容为中文，格式为角色一的姓名和人设，角色二的姓名和人设：\n{text}\n"
    print(content)
    response = client.chat.completions.create(
        model="glm-4",
        messages=[
            {
                "role": "user",
                "content": content
            }
        ],
        top_p=0.7,
        temperature=0.9,
        stream=False,
        max_tokens=2000,
    )
    # 提取角色和人设
    print(response.choices[0].message.content)
    text = response.choices[0].message.content
    text = text[7:]
    text = text[:-3]
    text = json.loads(text)
    return text

def save_file(text, file_name):
    with open(file_name, 'w', encoding='utf-8') as f:
        f.write(text)

if __name__ == "__main__":

    print("欢迎来到平凡的世界，请输入对话发生运行的次数(正整数)...\n")
    time = input()

    input_text = "贺秀莲是贺凤英的远门侄女，山西柳林人，从小丧母，跟随父亲长大。" \
                 "已经到了谈婚论嫁的年龄，别人也说了好几个亲事，但有自己想法的秀莲愣是谁都看不上。" \
                 "那天，家里来了个壮实的小伙子，看起来老实朴素，有着农民好手的精干强壮，秀莲第一眼便倾心了，这个人便是孙少安。" \
                 "前段时间还处在润叶撩起的纷乱思绪中，看到贺秀莲，这个汉子便抛去了所有杂念。" \
                 "眼前的人儿正是自己要寻找的那个——黑眉大眼、身材丰满、吃苦耐劳，关键是不要礼钱！" \
                 "这对孙少安这个贫困的人家来说是最吸引人的地方。"

    text = generate_role_play(input_text)
    bot_name = text["角色一"]["姓名"]
    print(bot_name)
    bot_info = text["角色一"]["人设"]
    print(bot_info)
    user_name = text["角色二"]["姓名"]
    print(user_name)
    user_info = text["角色二"]["人设"]
    print(user_info)

    messages = [
        {"role": "assistant", "content": "少安哥，我想去烫头"},
        {"role": "user", "content": "哎，秀莲，去吧"}
    ]

    # 生成time轮对话
    for i in range(int(time)):
        messages = characterglm_assistant(messages, bot_name, bot_info, user_name, user_info)
        messages = characterglm_user(messages, bot_name, bot_info, user_name, user_info)

    # print(messages)
    save_file(str(messages),"history.txt")
    print(time + "次对话已生成，保存在history.txt")

